from NN.CAE import *
from NN.DNN import *
import numpy as np
import random

def setup_seed(seed):
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
     np.random.seed(seed)
     random.seed(seed)
     torch.backends.cudnn.deterministic = True

def cifar_ten():
    # 设置随机数种子
    setup_seed(20)

    training_data = datasets.CIFAR10(
        root="data",
        train=True,
        download=True,
        transform=ToTensor(),
    )

    test_data = datasets.CIFAR10(
        root="data",
        train=False,
        download=True,
        transform=ToTensor()
    )
    t_len = int(len(training_data) * 0.7)
    v_len = len(training_data) - t_len
    train_db, val_db = torch.utils.data.random_split(training_data, [t_len, v_len])

    batch_size = 64
    train_dataloader = DataLoader(train_db, batch_size=batch_size, shuffle=True)
    val_dataloader = DataLoader(val_db, batch_size=batch_size, shuffle=True)
    test_dataloader = DataLoader(test_data, batch_size=batch_size)

    for X, y in test_dataloader:
        print("Shape of X [N, C, H, W]: ", X.shape)
        print("Shape of y: ", y.shape)
        break

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using {device} device")

    model_note = "Adam_e4"
    model = DNN(model_note=model_note, in_channels=3, output_shape=10, device_=device).to(device)
    # 加载上一次的保存的模型
    # model.pre_train(32)
    model.is_save_weights(True)
    # 显示实时训练结果
    model.set_draw_real_time()

    print(model)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.RMSprop(model.parameters(), lr=1e-4)

    epochs = 50
    for t in range(epochs):
        model.cae_train(train_dataloader, loss_fn, optimizer)
        model.cae_val(val_dataloader, loss_fn)
    model.cae_test(test_dataloader, loss_fn=loss_fn)
    print("Done!")
